package com.example.appjava;

import net.sf.jsqlparser.JSQLParserException;
import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.expression.LongValue;
import net.sf.jsqlparser.expression.Parenthesis;
import net.sf.jsqlparser.expression.operators.conditional.AndExpression;
import net.sf.jsqlparser.expression.operators.relational.ExpressionList;
import net.sf.jsqlparser.expression.operators.relational.ItemsList;
import net.sf.jsqlparser.expression.operators.relational.MultiExpressionList;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.schema.Column;
import net.sf.jsqlparser.schema.Table;
import net.sf.jsqlparser.statement.Statement;
import net.sf.jsqlparser.statement.Statements;
import net.sf.jsqlparser.statement.delete.Delete;
import net.sf.jsqlparser.statement.insert.Insert;
import net.sf.jsqlparser.statement.select.*;
import net.sf.jsqlparser.statement.update.Update;
import net.sf.jsqlparser.statement.values.ValuesStatement;
import org.hibernate.resource.jdbc.spi.StatementInspector;

import java.util.List;

/**
 * @author sjw
 * @className TenantStatementInspector
 * @date 2021/7/7 15:09
 * @description: TODO
 */
public class TenantStatementInspector implements StatementInspector {
    @Override
    public String inspect(String sql) {
        System.out.println("old sql: " + sql);
        try {
            Statements statements = CCJSqlParserUtil.parseStatements(sql);
            StringBuilder sqlStringBuilder = new StringBuilder();
            int i = 0;
            for (Statement statement : statements.getStatements()) {
                if (null != statement) {
                    if (i++ > 0) {
                        sqlStringBuilder.append(';');
                    }
                    sqlStringBuilder.append(this.processParser(statement));
                }
            }
            String newSql = sqlStringBuilder.toString();
            System.out.println("new sql: " + newSql);
            return newSql;
        } catch (JSQLParserException e) {
            e.printStackTrace();
        }
        return sql;
    }

    private String processParser(Statement statement) throws JSQLParserException {
        if (statement instanceof Insert) {
            this.processInsert((Insert) statement);
        } else if (statement instanceof Select) {
            this.processSelect(((Select) statement).getSelectBody());
        } else if (statement instanceof Update) {
            this.processUpdate((Update) statement);
        } else if (statement instanceof Delete) {
            this.processDelete((Delete) statement);
        }
        return statement.toString();
    }

    private void processInsert(Insert insert) {
        boolean exist = insert.getColumns().stream().anyMatch(e -> e.getColumnName().equals("tenant_id"));
        if (exist) return;
        insert.getColumns().add(new Column("tenant_id"));
        ItemsList itemsList = insert.getItemsList();
        if (itemsList instanceof ExpressionList) {
            ((ExpressionList) itemsList).getExpressions().add(new LongValue(1));
        } else if (itemsList instanceof MultiExpressionList) {
            ((MultiExpressionList) itemsList).getExprList().forEach(e -> {
                e.getExpressions().add(new LongValue(1));
            });
        }
    }

    private void processDelete(Delete delete) throws JSQLParserException {
        Expression oldWhere = delete.getWhere();
        Table table = delete.getTable();
        Expression where = CCJSqlParserUtil.parseCondExpression((table.getAlias() == null ? table.getName() : table.getAlias().getName()) + ".tenant_id=1");
        if (oldWhere != null)
            delete.setWhere(new AndExpression(new Parenthesis(oldWhere), where));
        else
            delete.setWhere(where);
    }

    private void processUpdate(Update update) throws JSQLParserException {
        Expression oldWhere = update.getWhere();
        Table table = update.getTable();
        Expression where = CCJSqlParserUtil.parseCondExpression((table.getAlias() == null ? table.getName() : table.getAlias().getName()) + ".tenant_id=1");
        if (oldWhere != null)
            update.setWhere(new AndExpression(new Parenthesis(oldWhere), where));
        else
            update.setWhere(where);

        List<Join> startJoin = update.getStartJoins();
        if (startJoin != null && startJoin.size() > 0) {
            for (Join join : startJoin) {
                processUpdateForStartJoin(join, update);
            }
        }
    }

    private void processUpdateForStartJoin(Join join, Update update) throws JSQLParserException {
        FromItem fromItem = join.getRightItem();
        if (fromItem instanceof Table) {
            processUpdateTable((Table) fromItem, update);
        } else if (fromItem instanceof SubSelect) {
            processUpdateSubSelect((SubSelect) fromItem, update);
        }
    }

    private void processUpdateTable(Table table, Update update) throws JSQLParserException {
        Expression oldWhere = update.getWhere();
        Expression where = CCJSqlParserUtil.parseCondExpression((table.getAlias() == null ? table.getName() : table.getAlias().getName()) + ".tenant_id=1");
        if (oldWhere != null)
            update.setWhere(new AndExpression(new Parenthesis(oldWhere), where));
        else
            update.setWhere(where);
    }

    private void processUpdateSubSelect(SubSelect subSelect, Update update) throws JSQLParserException {
        this.processSelect(subSelect.getSelectBody());
    }

    private void processPlainSelect(FromItem fromItem, PlainSelect plainSelect) throws JSQLParserException {
        if (fromItem instanceof Table) {
            processPlainTableSelect((Table) fromItem, plainSelect);
        } else if (fromItem instanceof SubSelect) {
            processPlainSubSelect((SubSelect) fromItem, plainSelect);
        }
    }

    private void processPlainSelect(PlainSelect plainSelect) throws JSQLParserException {
        FromItem fromItem = plainSelect.getFromItem();
        processPlainSelect(fromItem, plainSelect);
        List<Join> joins = plainSelect.getJoins();
        if (joins == null || joins.size() == 0) return;
        for (Join join : joins) {
            processPlainSelectJoin(join, plainSelect);
        }
    }

    private void processPlainSelectJoin(Join join, PlainSelect plainSelect) throws JSQLParserException {
        processPlainSelect(join.getRightItem(), plainSelect);
    }

    private void processPlainTableSelect(Table table, PlainSelect plainSelect) throws JSQLParserException {
        Expression oldWhere = plainSelect.getWhere();
        Expression where = CCJSqlParserUtil.parseCondExpression((table.getAlias() == null ? table.getName() : table.getAlias().getName()) + ".tenant_id=1");
        if (oldWhere != null)
            plainSelect.setWhere(new AndExpression(new Parenthesis(oldWhere), where));
        else
            plainSelect.setWhere(where);
    }

    private void processPlainSubSelect(SubSelect subSelect, PlainSelect plainSelect) throws JSQLParserException {
        this.processSelect(subSelect.getSelectBody());
    }

    private void processSelect(SelectBody selectBody) throws JSQLParserException {
        // The core of a "SELECT" statement (no UNION, no ORDER BY)
        if (selectBody instanceof PlainSelect) {
            processPlainSelect((PlainSelect) selectBody);
        }
        // One of the parts of a "WITH" clause of a "SELECT" statement
        else if (selectBody instanceof WithItem) {
            SelectBody withItemSelectBody = ((WithItem) selectBody).getSelectBody();
            if (withItemSelectBody != null) {
                processSelect(withItemSelectBody);
            }
        } else if (selectBody instanceof ValuesStatement) {
        }
        // A database set operation. This operation consists of a list of plainSelects connected by set operations (UNION,INTERSECT,MINUS,EXCEPT). All these operations have the same priority.
        else if (selectBody instanceof SetOperationList) {
            SetOperationList operationList = (SetOperationList) selectBody;
            if (operationList.getSelects() != null && operationList.getSelects().size() > 0) {
                for (SelectBody body : operationList.getSelects()) {
                    processSelect(body);
                }
            }
        }
    }
}
